package info.wangyuan.agent.service.rag;

import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import info.wangyuan.agent.common.Constant.ModelConstant;
import info.wangyuan.agent.common.Enum.ModelType;
import info.wangyuan.agent.entity.dto.RerankRequestDTO;
import info.wangyuan.agent.entity.dto.RerankResponseDTO;
import info.wangyuan.agent.entity.po.ModelConfig;
import info.wangyuan.agent.mapper.ModelConfigMapper;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

/**
 * 重排序处理
 *
 * @author Albert
 * @since 2025-09-07 16:39:14
 */
@Service
public class RerankDomainService {

    @Resource
    private RerankForestApi rerankForestApi;

    private final ModelConfigMapper modelConfigMapper;

    public RerankDomainService(ModelConfigMapper modelConfigMapper) {
        this.modelConfigMapper = modelConfigMapper;
    }

    public List<EmbeddingMatch<TextSegment>> rerankDocument(
            EmbeddingSearchResult<TextSegment> textSegmentEmbeddingSearchResult,
            String question
    ) {
        ModelConfig rerankModelConfig = modelConfigMapper.selectOne(
                Wrappers.lambdaQuery(ModelConfig.class)
                        .eq(ModelConfig::getType, ModelType.RERANK.getValue())
                        .eq(ModelConfig::getStatus, ModelConstant.IS_ACTIVE)
        );
        List<EmbeddingMatch<TextSegment>> matches = new ArrayList<>();

        final List<String> list = textSegmentEmbeddingSearchResult.matches().stream()
                .map(text -> text.embedded().text()).toList();

        final RerankRequestDTO rerankRequest = new RerankRequestDTO();
        rerankRequest.setModel(rerankModelConfig.getName());
        rerankRequest.setQuery(question);
        rerankRequest.setDocuments(list);

        // 使用Forest接口调用Rerank API
        final RerankResponseDTO rerankResponse = rerankForestApi.rerank(rerankModelConfig.getBaseUrl() + "/rerank",
                rerankModelConfig.getApiKey(), rerankRequest);

        final List<RerankResponseDTO.SearchResult> results = rerankResponse.getResults();

        results.forEach(result -> {
            final Integer index = result.getIndex();
            matches.add(textSegmentEmbeddingSearchResult.matches().get(index));
        });

        return matches;
    }

}
